import numpy as np
from dataclasses import dataclass

@dataclass
class State:
    position : np.ndarray
    velocity: np.ndarray
    euler    : np.ndarray
    angular_velocity: np.ndarray

class kinematics():
    """
    四旋翼运动学模型
    """
    def __init__(self,init_state: State = None):
        """
        初始化无人机状态
        """
        self.state = init_state if init_state else State(
            position = np.zeros(3,dtype=np.float64),
            velocity = np.zeros(3,dtype=np.float64),
            euler = np.zeros(3,dtype=np.float64),
            angular_velocity = np.zeros(3,dtype=np.float64)
        )
        self.quaternion = self.euler_to_quaternion(self.state.euler)
    
    def update_position(self,acceleration: np.ndarray, dt: float) -> None:
        """
        更新自身位置
        """
        self.state.position +=  self.state.velocity * dt
        self.state.velocity += acceleration * dt
    @staticmethod
    def euler_to_quaternion(euler : np.ndarray):
        """
        欧拉角转四元数 (旋转顺序: ZYX, 即yaw-pitch-roll)
        """
        roll, pitch, yaw = euler
        
        cy = np.cos(yaw * 0.5)
        sy = np.sin(yaw * 0.5)
        cr = np.cos(roll * 0.5)
        sr = np.sin(roll * 0.5)
        cp = np.cos(pitch * 0.5)
        sp = np.sin(pitch * 0.5)
        
        qw = cy * cr * cp + sy * sr * sp #检验过
        qx = cy * sr * cp - sy * cr * sp #检验过
        qy = cy * cr * sp + sy * sr * cp #检验过
        qz = sy * cr * cp - cy * sr * sp #检验过
        #返回float64
        return np.array([qw, qx, qy, qz], dtype=np.float64)
    
    @staticmethod
    def quaternion_to_euler(quaternion: np.ndarray):
        """
        将四元数转换为欧拉角 (旋转顺序: ZYX)
        """
        qw, qx, qy, qz = quaternion
        
        # 绕x轴旋转 (roll)
        sinr_cosp = 2 * (qw * qx + qy * qz)
        cosr_cosp = 1 - 2 * (qx**2 + qy**2)
        roll = np.arctan2(sinr_cosp, cosr_cosp)
        
        # 绕y轴旋转 (pitch)
        sinp = 2 * (qw * qy - qz * qx)
        if np.abs(sinp) >= 1:
            pitch = np.sign(sinp) * np.pi / 2  # 使用90度，如果超出范围
        else:
            pitch = np.arcsin(sinp)
        
        # 绕z轴旋转 (yaw)
        siny_cosp = 2 * (qw * qz + qx * qy)
        cosy_cosp = 1 - 2 * (qy**2 + qz**2)
        yaw = np.arctan2(siny_cosp, cosy_cosp)
        
        return np.array([roll, pitch, yaw])
    
    def update_attitude_quaternion(self, angular_acceleration: np.ndarray, dt: float) -> None:
        """
        使用四元数方法更新姿态 (避免万向节锁问题)
        """
        # 构建角速度四元数
        omega = self.state.angular_velocity
        q_omega = np.array([0, omega[0], omega[1], omega[2]], dtype=np.float64)
        
        # 计算四元数导数
        dq = 0.5 * self.quat_multiply(self.quaternion,q_omega)
        
        # 更新四元数
        self.quaternion += dq * dt
        
        # 归一化四元数
        self.quaternion = self.normalize_quaternion(self.quaternion)
        
        # 更新欧拉角状态
        self.state.euler = self.quaternion_to_euler(self.quaternion)

        self.state.angular_velocity += angular_acceleration * dt  # 假设角加速度是常数

    @staticmethod
    def quat_multiply(q1: np.ndarray, q2: np.ndarray) -> np.ndarray:
        """四元数乘法"""
        w1, x1, y1, z1 = q1
        w2, x2, y2, z2 = q2
        return np.array([
            w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, #检验过
            w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, #检验过
            w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, #检验过
            w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2  #检验过
        ], dtype=np.float64)
    
    @staticmethod
    def normalize_quaternion(quat: np.ndarray) -> np.ndarray:
        """归一化四元数"""
        norm = np.linalg.norm(quat)
        return quat / norm if norm > 0 else np.array([1, 0, 0, 0])
    
    def get_state(self):
        """返回当前状态，封装性"""
        return self.state
    
    def set_state(self,state: State):
        """设置变量,统一访问原则"""
        self.state = state
        self.quaternion = self.euler_to_quaternion(self.state.euler)
    
# if __name__ == "__main__":
#     init_state_xjr= State(np.array([0,0,0]),np.array([0,0,0]),np.array([0,0,0]),np.array([0,0,0]))
#     quad_kinematic_xjr = kinematics(init_state_xjr)
#     velocity_xjr = np.array([1,0,0])
#     angular_velocity = np.array([0,0,1])
#     dt = 0.1
#     for _ in range(10):
#         quad_kinematic_xjr.update_position(velocity_xjr,dt)
#         quad_kinematic_xjr.update_attitude_quaternion(angular_velocity,dt)
#         print(f"state:{quad_kinematic_xjr.get_state()}")
